import sys
import pylab as plt
import numpy as np
from params import *
import math
import os
import networkx as nx

datasets = ['bkite','fsquare','gowalla']
lbls = ['Brightkite', 'Foursquare', 'Gowalla']
node_dist = [5651.0,8494.0,5663.0]
link_dist = [2041.0,1442.0,1792.0]
degs = [7.88,22.07,9.48,7.0]


def scale_degree(d):
    M1 = 0.0
    M2 = 7.0
    MIN = 0
    MAX = 10**M2
    STEPS = 100
    l = math.log10(d)
    ratio = (l-M1)/(M2-M1)
    index = float(int(ratio*STEPS))/STEPS
    f = M1 + index*(M2-M1)
    #print d,l, ratio,index,f
    return 10**f

for dataset in datasets:
    trace_file = os.path.join(WORKDIR,'gsn','traces',dataset, '%s_graph.txt'%dataset)
    social_file = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,
            '%s_social_null_graph_1.txt'%dataset)
    geo_file = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,
            '%s_geo_null_graph_1.txt'%dataset)
    print 'Dataset %s'%dataset

    def plot_data(file):
        print file
        s = 0.0
        k = 0
        i = 0
        degrees = {}
        for line in open(file):
            i += 1
            if i % 1000000 == 0:
                print 'line ', i
            if line.startswith('#'):
                continue
            u1,u2,dist = line.split(' ')
            u1,u2 = map(int,(u1,u2))

            degrees.setdefault(u1,0)
            degrees[u1] += 1
            degrees.setdefault(u2,0)
            degrees[u2] += 1

            dist = float(dist)
            #dist = max((1.0,dist))
            s += dist
            k += 1

        avg_dist = s/k
        print "average link length", avg_dist
        print "links ", k

        res = {}
        for line in open(file):
            i += 1
            if i % 1000000 == 0:
                print 'line ', i
            if line.startswith('#'):
                continue
            u1,u2,dist = line.split(' ')
            u1,u2 = map(int,(u1,u2))
            dist = float(dist)

            k1 = degrees[u1]
            k2 = degrees[u2]
            kk = scale_degree(k1*k2)
            res.setdefault(kk,[]).append(dist)

        avg_dist=1
        x = sorted(res)
        y = [sum(res[k])/len(res[k])/avg_dist for k in x]

        return x,y

    lbls = []
    x,y = plot_data(trace_file)
    lbls.append('Original data')
    gx,gy = plot_data(geo_file)
    lbls.append('Geo model')
    sx,sy = plot_data(social_file)
    lbls.append('Social model')

    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    plt.loglog(x,y,'kx')
    plt.loglog(gx,gy,'k^', mfc='None')
    plt.loglog(sx,sy,'ko', mfc='None')
    plt.legend(lbls, loc='lower right',numpoints=1)
    plt.grid(True)
    plt.xlabel(r'$k_i k_j$')
    plt.ylabel(r'$\langle l_{ij} \rangle$')
    plt.axis((1,1e6,1e2,1e4))
    plt.savefig('%s_gravity.pdf'%dataset)
    plt.close()

